#!/usr/bin/env python
# coding=utf-8
# @FileName :ssh_tunnel.py
# @Time     :2023/4/10 16:46
# @Author   :ZXC
# @Desc     :ssh隧道

from loguru import logger
from sshtunnel import SSHTunnelForwarder

from .schemas import SSHTunnelObj


class SSHTunnel(object):
    """
    SSH隧道
    """

    def __init__(self, ssh_tunnel: SSHTunnelObj):
        """
        初始化并开启ssh隧道
        :param ssh_tunnel:
        """
        self._ssh_tunnel = ssh_tunnel
        logger.debug(f'正在初始化SSH隧道：{self._ssh_tunnel.dict()}')
        self._server = SSHTunnelForwarder(
            ssh_address_or_host=(self._ssh_tunnel.ssh_host, self._ssh_tunnel.ssh_port),
            ssh_username=self._ssh_tunnel.ssh_username,
            ssh_password=self._ssh_tunnel.ssh_password,
            remote_bind_address=(self._ssh_tunnel.remote_bind_host, self._ssh_tunnel.remote_bind_port),
        )
        self._server.start()
        logger.debug(f'SSH隧道初始化完成：{self._ssh_tunnel.dict()}')

    @property
    def server(self):
        """
        获取隧道
        :return:
        """
        return self._server

    def close(self):
        """
        关闭隧道
        :return:
        """
        logger.debug(f'正在关闭SSH隧道：{self._ssh_tunnel.dict()}')
        self._server.close()
        logger.debug(f'SSH隧道已关闭：{self._ssh_tunnel.dict()}')

